import sys
import cv2
import numpy as np
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QMainWindow, QFileDialog, QMessageBox, QLabel
from PyQt5.QtCore import Qt, pyqtSignal
from ui_interface import Ui_MainWindow
from functions import extract_scale_from_image, extract_scale_bar
import filetype
import re
from sam_functions import btnSAM_Clicked, sam_add_point, sam_segment, sam_undo_point
from measure import start_measurement
from save import save_results, enable_save_button


# 自定义可交互的QLabel，用于处理鼠标事件
class InteractiveLabel(QLabel):
    leftClicked = pyqtSignal(int, int)
    rightClicked = pyqtSignal(int, int)
    enterPressed = pyqtSignal()

    def __init__(self, parent=None):
        super(InteractiveLabel, self).__init__(parent)
        self.setMouseTracking(True)
        self.setFocusPolicy(Qt.StrongFocus)

    def mousePressEvent(self, event):
        pixmap = self.pixmap()
        if pixmap and hasattr(self.window(), 'sam_mode') and self.window().sam_mode:
            # 获取图像在标签中的实际显示区域
            scaled_width = pixmap.width()
            scaled_height = pixmap.height()

            # 计算图像在标签中的居中位置
            x_offset = (self.width() - scaled_width) / 2
            y_offset = (self.height() - scaled_height) / 2

            # 获取点击位置
            pos_x = event.x() - x_offset
            pos_y = event.y() - y_offset

            # 判断点击是否在图像区域内
            if 0 <= pos_x < scaled_width and 0 <= pos_y < scaled_height:
                # 根据鼠标按钮发送不同信号
                if event.button() == Qt.LeftButton:
                    self.leftClicked.emit(int(pos_x), int(pos_y))
                elif event.button() == Qt.RightButton:
                    self.rightClicked.emit(int(pos_x), int(pos_y))

        super(InteractiveLabel, self).mousePressEvent(event)

    def keyPressEvent(self, event):
        if event.key() == Qt.Key_Return or event.key() == Qt.Key_Enter:
            self.enterPressed.emit()
        super(InteractiveLabel, self).keyPressEvent(event)


class InsectMeasureApp(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super(InsectMeasureApp, self).__init__()
        self.setupUi(self)
        self.setWindowTitle("图像测量工具")

        # 初始化变量
        self.source_image = None
        self.drawing_image = None
        self.displayed_image = None
        self.segmented_image = None
        self.mask = None
        self.scale_ratio = 1.0
        self.sam_mode = False
        self.input_points = []
        self.input_labels = []

        # 添加撤销相关变量
        self.input_points_history = []
        self.input_labels_history = []
        self.drawing_image_history = []

        # 替换原有的QLabel为自定义的InteractiveLabel，但保留原来的位置和尺寸
        self.setup_interactive_label()

        # 连接信号和槽
        self.actionOpen.triggered.connect(self.open_image)
        self.btnOpen.clicked.connect(self.open_image)

        # 连接比例尺按钮
        self.btnSegment.clicked.connect(self.btnScale_Clicked)

        # 连接SAM按钮
        self.btnSAM.clicked.connect(self.btnSAM_Clicked)

        # 连接测量按钮
        self.btnMeasure.clicked.connect(lambda: start_measurement(self))

        # 连接保存按钮
        self.btnSave.clicked.connect(lambda: save_results(self))
        self.actionSave.triggered.connect(lambda: save_results(self))

        # 绑定撤销快捷键 - 使用Ctrl+Z
        self.undo_shortcut = QtWidgets.QShortcut(QtGui.QKeySequence("Ctrl+Z"), self)
        self.undo_shortcut.activated.connect(self.sam_undo_point)

        # 暂时禁用未实现的功能按钮
        self.btnSegment.setEnabled(True)  # 启用比例尺按钮
        self.btnMeasure.setEnabled(False)
        self.btnSAM.setEnabled(True)  # 启用SAM按钮
        self.btnErase.setEnabled(False)
        self.btnSave.setEnabled(False)
        self.actionSave.setEnabled(False)

        # 设置窗口大小调整计时器
        self.resizeTimer = QtCore.QTimer()
        self.resizeTimer.setSingleShot(True)
        self.resizeTimer.timeout.connect(self.on_resize_timeout)

        # 设置状态栏初始消息
        self.statusbar.showMessage("准备就绪。请打开一个图像文件。")

    def resizeEvent(self, event):
        """处理窗口大小调整事件，包括最大化/全屏"""
        super(InsectMeasureApp, self).resizeEvent(event)

        # 立即更新UI布局，确保获取到最新的尺寸
        QtWidgets.QApplication.processEvents()

        # 无需等待，立即调整图像大小
        self.on_resize_timeout()

        # 同时启动延时器，以应对连续的调整
        self.resizeTimer.start(200)

    def on_resize_timeout(self):
        """当窗口大小调整完成后更新所有图像"""
        if self.source_image is None:
            return

        # 强制刷新布局，确保获取最新尺寸
        self.centralwidget.updateGeometry()
        self.lblSourceImage.updateGeometry()
        self.lblSegmentedImage.updateGeometry()
        self.lblSingleObject.updateGeometry()

        # 刷新应用程序，确保布局已重新计算
        QtWidgets.QApplication.processEvents()

        # 确定并重新显示图像
        display_img = self.drawing_image if (hasattr(self, 'drawing_image') and
                                             self.drawing_image is not None and
                                             self.sam_mode) else self.source_image

        # 重新计算并显示原始图像
        self.display_image(display_img)

        # 更新分割图像和单个对象图像(如果存在)
        if hasattr(self, 'segmented_image') and self.segmented_image is not None:
            self.display_segmented_image(self.segmented_image)

        if hasattr(self, 'mask') and self.mask is not None:
            self.display_single_object(self.mask)

    def setup_interactive_label(self):
        """设置交互式标签，保留原来的布局和属性"""
        # 获取原标签的位置和尺寸信息
        original_label = self.lblSourceImage
        geo = original_label.geometry()
        parent = original_label.parent()
        name = original_label.objectName()
        style = original_label.styleSheet()
        layout_item = None

        # 如果原标签在布局中，获取布局信息
        if parent.layout() is not None:
            layout = parent.layout()
            index = -1
            for i in range(layout.count()):
                if layout.itemAt(i).widget() == original_label:
                    index = i
                    layout_item = layout.itemAt(i)
                    break

            if index != -1:
                layout.removeItem(layout_item)

        # 创建新的交互式标签
        self.lblSourceImage = InteractiveLabel(parent)
        self.lblSourceImage.setObjectName(name)
        self.lblSourceImage.setStyleSheet(style)
        self.lblSourceImage.setGeometry(geo)
        self.lblSourceImage.setAlignment(Qt.AlignCenter)

        # 确保标签可以调整大小
        self.lblSourceImage.setSizePolicy(QtWidgets.QSizePolicy.Expanding,
                                          QtWidgets.QSizePolicy.Expanding)
        self.lblSourceImage.setMinimumSize(300, 300)  # 设置最小尺寸

        # 如果原来在布局中，添加到相同位置
        if layout_item is not None and parent.layout() is not None:
            parent.layout().addWidget(self.lblSourceImage)

        # 连接信号
        self.lblSourceImage.leftClicked.connect(lambda x, y: self.sam_add_point(x, y, True))
        self.lblSourceImage.rightClicked.connect(lambda x, y: self.sam_add_point(x, y, False))
        self.lblSourceImage.enterPressed.connect(self.sam_segment)

        # 删除原标签
        original_label.deleteLater()

    def open_image(self):
        """从本地读取图片并提取比例尺信息"""
        # 重置变量
        self.source_image = None
        self.drawing_image = None
        self.displayed_image = None
        self.segmented_image = None
        self.mask = None
        self.scale_ratio = 1.0
        self.input_points_history = []
        self.input_labels_history = []
        self.drawing_image_history = []

        # 打开文件选择对话框
        filename, _ = QFileDialog.getOpenFileName(
            self,
            "打开图像",
            "",  # 可以根据需要设置默认路径
            "图像文件 (*.png *.jpg *.jpeg *.bmp *.tif)"
        )

        if not filename:
            return  # 用户取消了操作

        try:
            # 检查文件类型
            kind = filetype.guess(filename)
            if kind is None:
                QMessageBox.warning(self, "警告", "无法识别的文件类型")
                return

            # 读取图像
            self.source_image = cv2.imread(filename)
            if self.source_image is None:
                QMessageBox.warning(self, "警告", "无法读取图像，请检查文件格式或路径")
                return

            # 创建图像副本用于绘制
            self.drawing_image = self.source_image.copy()

            # 显示图像
            self.display_image()

            # 提取比例尺信息
            self.extract_scale_bar_info()

            # 更新状态栏
            self.statusbar.showMessage(f"已加载图像: {filename}")

            # 启用相关按钮
            self.btnSegment.setEnabled(True)
            self.btnSAM.setEnabled(True)

            # 清空分割和对象图像框
            self.lblSegmentedImage.clear()
            self.lblSegmentedImage.setText("分割图像")
            self.lblSingleObject.clear()
            self.lblSingleObject.setText("目标对象")

        except Exception as e:
            QMessageBox.critical(self, "错误", f"打开图像时发生错误: {str(e)}")
            self.statusbar.showMessage("打开图像失败")

    def extract_scale_bar_info(self):
        """从图像中提取比例尺信息并更新UI"""
        if self.source_image is None:
            return

        # 调用移动到functions.py的函数
        scale_pixel_length, scale_value, scale_unit = extract_scale_bar(self.source_image)

        if scale_pixel_length is not None:
            self.txtUnitLength.setText(str(scale_pixel_length))

            if scale_value is not None:
                self.txtUnitPixels.setText(scale_value)

                # 设置单位选择
                if scale_unit in ['μm', 'um']:
                    self.radioUM.setChecked(True)
                elif scale_unit == 'mm':
                    self.radioMM.setChecked(True)
                elif scale_unit == 'm':
                    self.radioM.setChecked(True)

                self.statusbar.showMessage(
                    f"已提取比例尺: {scale_value} {scale_unit}, 像素长度: {scale_pixel_length}")
            else:
                self.statusbar.showMessage(f"已提取比例尺像素长度: {scale_pixel_length}")
        else:
            self.statusbar.showMessage("未找到比例尺")

    def display_image(self, image=None):
        """根据标签大小调整并显示图像"""
        # 如果没有指定图像，使用source_image
        display_img = image if image is not None else self.source_image

        if display_img is None:
            return

        # 强制重新获取标签当前的实际尺寸
        QtWidgets.QApplication.processEvents()
        label_width = self.lblSourceImage.width()
        label_height = self.lblSourceImage.height()

        # 获取图像尺寸
        img_height, img_width, _ = display_img.shape

        # 计算缩放比例保持宽高比
        width_ratio = label_width / img_width
        height_ratio = label_height / img_height
        self.scale_ratio = min(width_ratio, height_ratio)

        # 计算调整后的尺寸
        new_width = int(img_width * self.scale_ratio)
        new_height = int(img_height * self.scale_ratio)

        # 调整图像大小
        resized_image = cv2.resize(display_img, (new_width, new_height))

        # 转换为QImage
        height, width, channel = resized_image.shape
        bytes_per_line = 3 * width
        q_img = QtGui.QImage(
            resized_image.data, width, height, bytes_per_line,
            QtGui.QImage.Format_RGB888
        ).rgbSwapped()

        # 使用QPixmap显示图像
        pixmap = QtGui.QPixmap.fromImage(q_img)
        self.lblSourceImage.setPixmap(pixmap)

        # 保存显示的图像
        self.displayed_image = resized_image

    def btnScale_Clicked(self):
        '''
        比例尺单击按钮
        :return:
        '''
        if self.source_image is None:
            QMessageBox.warning(self, "警告", "请先打开图像")
            return

        # 提取比例尺信息
        try:
            scale = extract_scale_from_image(self.source_image)  # 使用正确的变量名source_image
            self.statusbar.showMessage(f"提取到的比例尺: {scale}")

            # 解析比例尺信息
            match = re.match(r"([\d.]+)\s*(μm|um|mm|m)", scale, re.IGNORECASE)
            if match:
                value, unit = match.groups()
                self.txtUnitPixels.setText(value)  # 更新输入框的值

                # 更新单位选择
                if unit.lower() in ["μm", "um"]:
                    self.radioUM.setChecked(True)
                elif unit.lower() == "mm":
                    self.radioMM.setChecked(True)
                elif unit.lower() == "m":
                    self.radioM.setChecked(True)

                QMessageBox.information(self, "成功", f"检测到比例尺：{value} {unit}")
            else:
                self.statusbar.showMessage("无法解析比例尺信息")
                QMessageBox.warning(self, "警告", "无法解析比例尺信息")

        except Exception as e:
            self.statusbar.showMessage(f"提取比例尺时发生错误: {str(e)}")
            QMessageBox.critical(self, "错误", f"提取比例尺时发生错误: {str(e)}")

    # 使用sam_functions.py中的函数
    def btnSAM_Clicked(self):
        btnSAM_Clicked(self)

    def sam_add_point(self, x, y, is_positive=True):
        sam_add_point(self, x, y, is_positive)

    def sam_segment(self):
        sam_segment(self)
        # 成功分割后启用测量按钮
        self.btnMeasure.setEnabled(True)

    def sam_undo_point(self):
        sam_undo_point(self)

    def display_segmented_image(self, segmented_image):
        """显示分割图像"""
        if segmented_image is None:
            self.lblSegmentedImage.clear()
            self.lblSegmentedImage.setText("分割图像")
            return

        # 保存分割图像供窗口调整时使用
        self.segmented_image = segmented_image

        # 强制重新获取标签当前的实际尺寸
        QtWidgets.QApplication.processEvents()
        label_width = self.lblSegmentedImage.width()
        label_height = self.lblSegmentedImage.height()

        # 获取图像尺寸
        img_height, img_width, _ = segmented_image.shape

        # 计算缩放比例
        width_ratio = label_width / img_width
        height_ratio = label_height / img_height
        scale_ratio = min(width_ratio, height_ratio)

        # 计算调整后的尺寸
        new_width = int(img_width * scale_ratio)
        new_height = int(img_height * scale_ratio)

        # 调整图像大小
        resized_image = cv2.resize(segmented_image, (new_width, new_height))

        # 转换为QImage以显示
        height, width, channel = resized_image.shape
        bytes_per_line = 3 * width
        q_img = QtGui.QImage(
            resized_image.data, width, height, bytes_per_line,
            QtGui.QImage.Format_RGB888
        ).rgbSwapped()

        # 设置图像到标签
        self.lblSegmentedImage.setPixmap(QtGui.QPixmap.fromImage(q_img))

    def display_single_object(self, mask):
        """显示分割出的单个对象"""
        if self.source_image is None or mask is None:
            self.lblSingleObject.clear()
            self.lblSingleObject.setText("目标对象")
            return

        # 保存掩码供窗口调整时使用
        self.mask = mask

        # 创建一个与原始图像大小相同的空白图像
        single_object_image = np.zeros_like(self.source_image)

        # 使用mask提取分割出的对象
        single_object_image[mask] = self.source_image[mask]

        # 强制重新获取标签当前的实际尺寸
        QtWidgets.QApplication.processEvents()
        label_width = self.lblSingleObject.width()
        label_height = self.lblSingleObject.height()

        # 获取图像尺寸
        img_height, img_width, _ = single_object_image.shape

        # 计算缩放比例
        width_ratio = label_width / img_width
        height_ratio = label_height / img_height
        scale_ratio = min(width_ratio, height_ratio)

        # 计算调整后的尺寸
        new_width = int(img_width * scale_ratio)
        new_height = int(img_height * scale_ratio)

        # 调整图像大小
        resized_image = cv2.resize(single_object_image, (new_width, new_height))

        # 转换为QImage以显示
        height, width, channel = resized_image.shape
        bytes_per_line = 3 * width
        q_img = QtGui.QImage(
            resized_image.data, width, height, bytes_per_line,
            QtGui.QImage.Format_RGB888
        ).rgbSwapped()

        # 设置图像到标签
        self.lblSingleObject.setPixmap(QtGui.QPixmap.fromImage(q_img))


if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    window = InsectMeasureApp()
    window.show()
    sys.exit(app.exec_())
